import numpy as np
import pyarma as pa

def cd1(M, h):

    H = pa.ones(M, M)

    return (pa.diagmat(H, 1) - pa.diagmat(H, -1)) / (2 * h)

def cd2(M, h):

    H = pa.ones(M, M)
    
    return (pa.diagmat(H, 1) - 2 * pa.diagmat(H) + pa.diagmat(H, -1)) / (h ** 2)

def cd1x(Mx, My, hx, hy):

    D = pa.zeros(Mx * My, Mx * My)

    for i in range(Mx):
        D[i*My:(i+1)*My-1, i*My:(i+1)*My-1] = cd1(Mx, hx)

    return D

def cd1y(Mx, My, hx, hy):

    H = pa.ones(Mx * My, Mx * My)

    return (pa.diagmat(H, My) - pa.diagmat(H, -My)) / (hy ** 2)

def cd2x(Mx, My, hx, hy):

    D = pa.zeros(Mx * My, Mx * My)

    for i in range(Mx):
        D[i*My:(i+1)*My-1, i*My:(i+1)*My-1] = cd2(Mx, hx)

    return D

def cd2y(Mx, My, hx, hy):

    H = pa.ones(Mx * My, Mx * My)

    return (pa.diagmat(H, My) - 2 * pa.diagmat(H) + pa.diagmat(H, -My)) / (hy ** 2)

def GLK1010(M, h):

    H = pa.ones(M, M)
    
    return (pa.diagmat(H, 1) + 4 * pa.diagmat(H) + pa.diagmat(H, -1)) * h / 6

def GLK1111(M, h):

    H = pa.ones(M, M)
    
    return (- pa.diagmat(H, 1) + 2 * pa.diagmat(H) - pa.diagmat(H, -1)) / h

def GLK1011(M, h):

    H = pa.ones(M, M)
    
    return (-pa.diagmat(H, 1) + pa.diagmat(H, -1)) / 2

def GLK2011(M, h, u):
    
    un = pa.zeros(M, 1)

    for i in range(M):

        if i > 0:

            un[i] += u[i - 1] * u[i - 1] / 3 + u[i - 1] * u[i] / 3

        if i < M - 1:

            un[i] -= u[i + 1] * u[i + 1] / 3 + u[i + 1] * u[i] / 3

    return un

if __name__ == '__main__':

    print(cd1y(2, 3, 0.5, 0.5))

